//
// Created by tian on 2020/5/24.
//
#include "cmw_app/detection_app.h"
#include "common/project_root.h"

using namespace Yolo;
using namespace Tn;
using namespace std;
using namespace cv;

namespace adc{
DetectionApp::DetectionApp(const std::string &config_file, const float hz)  : Worker(hz)
{
    InitConfig(config_file);
    detection_callback_ = NULL;
}

void DetectionApp::SetDetectionCallback(DetectionCallback callback)
{
    detection_callback_ = callback;
}

vector<string> DetectionApp::GetNamesFromFile(string filePathAndName)
{
    vector<string> retVec;
    ifstream cocoName;
    cocoName.open(filePathAndName.data());
    assert(cocoName.is_open());
    string s;
    while (getline(cocoName,s))
    {
        retVec.push_back(s);
    }
    cocoName.close();
    return retVec;
}


void DetectionApp::InitConfig(const std::string &config_file)
{
    try
    {
        YAML::Node node = YAML::LoadFile(config_file);
        std::string class_name_path = adc::RootPath::GetAbsolutePath( node["class_name_path"].as<std::string>());
        std::string engine_path = adc::RootPath::GetAbsolutePath(node["engine_path"].as<std::string>());

        width_ = node["width"].as<int>(608);
	    hight_ = node["hight"].as<int>(608);
        channel_ = node["channel"].as<int>(3);
        nms_ = node["nms"].as<float>(0.5);
        net_.reset(new Tn::trtNet(engine_path));
        outputs_ = std::make_shared<std::list<vector<Tn::Bbox>>>();
        class_name_ = GetNamesFromFile(class_name_path);
    }
    catch(const std::exception& e)
    {
        std::cerr << e.what() << '\n';
    }
}

void DetectionApp::CacheImages(const cv::Mat &left_img, std::chrono::system_clock::time_point &time)
{
    std::lock_guard<std::mutex> mutex_cache_image(mutex_image_);

    image_cache_.push(left_img);
    time_cache_.push(time);
    while (image_cache_.size() > kQueueSize)
    {
        image_cache_.pop();
        time_cache_.pop();
    }
}

bool DetectionApp::ReceiveImage(void)
{
    std::lock_guard<std::mutex> mutex_cache_image(mutex_image_);
    if(image_cache_.empty())
    {
        return false;
    }
    left_image_ = image_cache_.front().clone();
    current_time_ = time_cache_.front();
    image_cache_.pop();
    time_cache_.pop();
    return true;
}

vector<float> DetectionApp::prepareImage(cv::Mat& img)
{
    using namespace cv;

    int c = channel_;
    int h = hight_;   //net h
    int w = width_;   //net w

    float scale = min(float(w)/img.cols,float(h)/img.rows);
    auto scaleSize = cv::Size(img.cols * scale,img.rows * scale);

    cv::Mat rgb ;
    cv::cvtColor(img, rgb, cv::COLOR_BGR2RGB);
    cv::Mat resized;
    cv::resize(rgb, resized,scaleSize,0,0,INTER_CUBIC);

    cv::Mat cropped(h, w,CV_8UC3, 127);
    Rect rect((w- scaleSize.width)/2, (h-scaleSize.height)/2, scaleSize.width,scaleSize.height);
    resized.copyTo(cropped(rect));

    cv::Mat img_float;
    if (c == 3)
        cropped.convertTo(img_float, CV_32FC3, 1/255.0);
    else
        cropped.convertTo(img_float, CV_32FC1 ,1/255.0);

    //HWC TO CHW
    vector<Mat> input_channels(c);
    cv::split(img_float, input_channels);

    vector<float> result(h*w*c);
    auto data = result.data();
    int channelLength = h * w;
    for (int i = 0; i < c; ++i) {
        memcpy(data,input_channels[i].data,channelLength*sizeof(float));
        data += channelLength;
    }

    return result;
}


void DetectionApp::DoNms(vector<Detection>& detections,int classes ,float nmsThresh)
{
    vector<vector<Detection>> resClass;
    resClass.resize(classes);

    for (const auto& item : detections)
        resClass[item.classId].push_back(item);

    auto iouCompute = [](float * lbox, float* rbox)
    {
        float interBox[] = {
                max(lbox[0] - lbox[2]/2.f , rbox[0] - rbox[2]/2.f), //left
                min(lbox[0] + lbox[2]/2.f , rbox[0] + rbox[2]/2.f), //right
                max(lbox[1] - lbox[3]/2.f , rbox[1] - rbox[3]/2.f), //top
                min(lbox[1] + lbox[3]/2.f , rbox[1] + rbox[3]/2.f), //bottom
        };

        if(interBox[2] > interBox[3] || interBox[0] > interBox[1])
            return 0.0f;

        float interBoxS =(interBox[1]-interBox[0])*(interBox[3]-interBox[2]);
        return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] -interBoxS);
    };

    vector<Detection> result;
    for (int i = 0;i<classes;++i)
    {
        auto& dets =resClass[i];
        if(dets.size() == 0)
            continue;

        sort(dets.begin(),dets.end(),[=](const Detection& left,const Detection& right){
            return left.prob > right.prob;
        });

        for (unsigned int m = 0;m < dets.size() ; ++m)
        {
            auto& item = dets[m];
            result.push_back(item);
            for(unsigned int n = m + 1;n < dets.size() ; ++n)
            {
                if (iouCompute(item.bbox,dets[n].bbox) > nmsThresh)
                {
                    dets.erase(dets.begin()+n);
                    --n;
                }
            }
        }
    }
    //swap(detections,result);
    detections = move(result);
}

vector<Bbox> DetectionApp::postProcessImg(cv::Mat& img,vector<Detection>& detections,int classes)
{
    using namespace cv;

    int h = hight_;   //net h
    int w = width_;   //net w

    //scale bbox to img
    int width = img.cols;
    int height = img.rows;
    float scale = min(float(w)/width,float(h)/height);
    float scaleSize[] = {width * scale,height * scale};

    //correct box
    for (auto& item : detections)
    {
        auto& bbox = item.bbox;
        bbox[0] = (bbox[0] * w - (w - scaleSize[0])/2.f) / scaleSize[0];
        bbox[1] = (bbox[1] * h - (h - scaleSize[1])/2.f) / scaleSize[1];
        bbox[2] /= scaleSize[0];
        bbox[3] /= scaleSize[1];
    }

    //nms
    float nmsThresh = nms_;
    if(nmsThresh > 0)
        DoNms(detections,classes,nmsThresh);

    vector<Bbox> boxes;
    for(const auto& item : detections)
    {
        auto& b = item.bbox;
        Bbox bbox =
                {
                        item.classId,   //classId
                        max(int((b[0]-b[2]/2.)*width),0), //left
                        min(int((b[0]+b[2]/2.)*width),width), //right
                        max(int((b[1]-b[3]/2.)*height),0), //top
                        min(int((b[1]+b[3]/2.)*height),height), //bot
                        item.prob       //score
                };
        boxes.push_back(bbox);
    }

    return boxes;
}

void DetectionApp::Inference()
{
    vector<float> curInput = prepareImage(left_image_);
    //Tn::RUN_MODE run_mode = Tn::RUN_MODE::FLOAT16;

    int outputCount = net_->getOutputSize()/sizeof(float);
    int classNum = 80;
    output_data_.reset(new float[outputCount]);
    net_->doInference(curInput.data(), output_data_.get(),1);
    auto output = output_data_.get();
    int detCount = output[0];
    vector<Detection>  result;

    result.resize(detCount);
    memcpy(result.data(), &output[1],detCount*sizeof(Detection));

    outputs_->clear();
    auto boxes = postProcessImg(left_image_,result,classNum);
    outputs_->emplace_back(boxes);
}




void DetectionApp::DoWork()
{
    static int count = 0;
    if(!ReceiveImage())
    {
        if(count>50)
        {
            std::cout<<"can not get images"<<std::endl;
        }
        count ++;
        return;
    }
    count = 0;
    
    Inference();


    if(NULL != detection_callback_)
    {
        detection_callback_(outputs_, current_time_,left_image_,class_name_);
    }
}
}



